#include <netinet/in.h>
#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <unistd.h>
#include <pthread.h>
#include <sched.h>
#include <string.h>
#include "sr_arpcache.h"
#include "sr_router.h"
#include "sr_icmp.h"
#include "sr_if.h"
#include "sr_protocol.h"
#include "sr_utils.h"

/* This function gets called every second. See the comments in the header file
   for an idea of what it should look like. */
void sr_arpcache_sweepreqs(struct sr_instance *sr) {
    fprintf(stderr, "SweepReqs called\n");
    // sr_arpcache_dump(&sr->cache);
    
    struct sr_arpreq *current_request = sr->cache.requests;
    fprintf(stderr, "current request %p\n", current_request);
    fprintf(stderr, "ETHER_ARP_PACKET_LEN: %ld", ETHER_ARP_PACKET_LEN);
    while (current_request)
    {
        sr_handle_arpreq(sr, current_request);
        fprintf(stderr, "requests %p\n", current_request);
        
        current_request = current_request->next;
    }
}

void sr_handle_arpreq(struct sr_instance *sr, struct sr_arpreq *arpreq) 
{
    if (difftime(time(NULL), arpreq->sent) > SR_ARPREQ_TIMEOUT)
    {
        if (arpreq->times_sent >= SR_ARPREQ_SEND)
        {   
            // Send IMCP host unreachable to all packets in arpreq
            // call destroy on arpreq
            struct sr_packet *waiting_packet = arpreq->packets;
            while (waiting_packet)
            {
                unsigned int icmp_pkt_len = sizeof(sr_ethernet_hdr_t) + 2*sizeof(sr_ip_hdr_t) + 
                                            sizeof(sr_icmp_hdr_t) + 8;
                uint8_t *icmp_unreachable_packet = malloc(icmp_pkt_len);
                struct sr_if *interface = sr_get_interface(sr, waiting_packet->iface);
                
                sr_ethernet_hdr_t *packet_ether_hdr = (sr_ethernet_hdr_t *)waiting_packet->buf;
                sr_ip_hdr_t *packet_ip_hdr = (sr_ip_hdr_t *)(waiting_packet->buf + sizeof(sr_ethernet_hdr_t));

                sr_ethernet_hdr_t *ether_hdr = (sr_ethernet_hdr_t *)icmp_unreachable_packet;
                memcpy(ether_hdr->ether_dhost, packet_ether_hdr->ether_shost, ETHER_ADDR_LEN);
                memcpy(ether_hdr->ether_shost, interface->addr, ETHER_ADDR_LEN);
                ether_hdr->ether_type = ntohs(ETHERTYPE_IP);

                sr_ip_hdr_t *outer_ip_hdr = (sr_ip_hdr_t *)(icmp_unreachable_packet + sizeof(sr_ethernet_hdr_t));
                outer_ip_hdr->ip_tos = 0;
                outer_ip_hdr->ip_len = ntohs(icmp_pkt_len - sizeof(sr_ethernet_hdr_t));
                outer_ip_hdr->ip_id = 0;
                outer_ip_hdr->ip_off = 0;
                outer_ip_hdr->ip_ttl = ICMP_TTL;
                outer_ip_hdr->ip_sum = 0; // Must zero out cksum first
                outer_ip_hdr->ip_p = IPPROTO_ICMP;
                outer_ip_hdr->ip_src = ntohl(interface->ip);
                outer_ip_hdr->ip_dst = packet_ip_hdr->ip_src; // but is it already in network order?
                // cksum already in network byte order
                outer_ip_hdr->ip_sum = cksum(outer_ip_hdr, sizeof(sr_ip_hdr_t)); 
                
                sr_icmp_hdr_t *icmp_hdr = (sr_icmp_hdr_t *)(icmp_unreachable_packet + sizeof(sr_ethernet_hdr_t) + sizeof(sr_ip_hdr_t));
                icmp_hdr->icmp_type = ICMP_UNREACHABLE_TYPE;
                icmp_hdr->icmp_code = ICMP_HOST_UNREACHABLE_CODE;
                icmp_hdr->icmp_sum = 0; // cksum should including ICMP payload
                icmp_hdr->icmp_sum = cksum(icmp_hdr, sizeof(sr_icmp_hdr_t));    
                
                // TODO copy packet
                
                waiting_packet = waiting_packet->next;
            }
            sr_arpreq_destroy(&sr->cache, arpreq);
        }
        else // Send ARP request packet
        {
            // Target should use first packet's interface's MAC and IP addresses
            struct sr_if *interface = sr_get_interface(sr, arpreq->packets->iface);
            uint8_t *arpreq_packet = sr_new_ether_arp_pkt(NULL,
                                                          interface->addr,
                                                          arpreq->ip,
                                                          interface->ip,
                                                          1);
            sr_send_packet(sr, arpreq_packet, ETHER_ARP_PACKET_LEN, interface->name);
            arpreq->sent = time(NULL);
            arpreq->times_sent++;
            
            free(arpreq_packet);
        }
    }
}

uint8_t *sr_new_ether_arp_pkt(const unsigned char *destination_MAC,
                              const unsigned char *sender_MAC,
                              uint32_t destination_ip,
                              uint32_t sender_ip,
                              int is_request)
{
    unsigned int arp_pkt_len = sizeof(sr_ethernet_hdr_t) + sizeof(sr_arp_hdr_t);
    uint8_t *arp_pkt = malloc(arp_pkt_len);
    
    // Cast packet into ethernet and arp headers for easy access
    sr_ethernet_hdr_t *ether_hdr = (sr_ethernet_hdr_t *) arp_pkt;
    sr_arp_hdr_t *arp_hdr = (sr_arp_hdr_t *) (arp_pkt + sizeof(sr_ethernet_hdr_t));

    if (destination_MAC)
    {
        memcpy(ether_hdr->ether_dhost, destination_MAC, ETHER_ADDR_LEN);
        memcpy(arp_hdr->ar_tha, destination_MAC, ETHER_ADDR_LEN);
    }
    else // broadcast ethernet packet, dhost: ff:ff:ff:ff:ff:ff, ar_tha: 00:..:00
    {
        memcpy(ether_hdr->ether_dhost, BROADCAST_MAC_ADDRESS, ETHER_ADDR_LEN);
        bzero(arp_hdr->ar_tha, ETHER_ADDR_LEN);
    }
    memcpy(ether_hdr->ether_shost, sender_MAC, ETHER_ADDR_LEN);
    ether_hdr->ether_type = ETHERTYPE_ARP;
    
    // Init rest of ARP data
    arp_hdr->ar_op = is_request? htons(ARP_REQUEST) : htons(ARP_REPLY); // ushort
    arp_hdr->ar_hrd = htons(ARP_HDR_ETHER);     // ushort
    arp_hdr->ar_pro = htons(ETHERTYPE_IP);      // ushort
    arp_hdr->ar_hln = ETHER_ADDR_LEN;           // uchar
    arp_hdr->ar_pln = sizeof(uint32_t);         // uchar
    
    memcpy(arp_hdr->ar_sha, sender_MAC, ETHER_ADDR_LEN);
    arp_hdr->ar_sip = htonl(sender_ip);         // ulong
    arp_hdr->ar_tip = htonl(destination_ip);    // ulong
    
    return arp_pkt;
}

/* You should not need to touch the rest of this code. */

/* Checks if an IP->MAC mapping is in the cache. IP is in network byte order. 
   You must free the returned structure if it is not NULL. */
struct sr_arpentry *sr_arpcache_lookup(struct sr_arpcache *cache, uint32_t ip) {
    pthread_mutex_lock(&(cache->lock));
    
    struct sr_arpentry *entry = NULL, *copy = NULL;
    
    int i;
    for (i = 0; i < SR_ARPCACHE_SZ; i++) {
        if ((cache->entries[i].valid) && (cache->entries[i].ip == ip)) {
            entry = &(cache->entries[i]);
        }
    }
    
    /* Must return a copy b/c another thread could jump in and modify
       table after we return. */
    if (entry) {
        copy = (struct sr_arpentry *) malloc(sizeof(struct sr_arpentry));
        memcpy(copy, entry, sizeof(struct sr_arpentry));
    }
        
    pthread_mutex_unlock(&(cache->lock));
    
    return copy;
}

/**
 * Handles external ARP request. Checks own if_list for an IP match and
 * sends out an ARP reply if there is a match.
 * TODO not sure what needs to be converted from ntoh
 */
void sr_handle_received_arpreq(struct sr_instance* sr, 
                               sr_arp_hdr_t *arp_header)
{
    uint32_t target_ip = ntohl(arp_header->ar_tip);
    unsigned char *target_hw_addr;
    int found_flag = 0;
    struct sr_if *current_if = sr->if_list;
    
    while (current_if)
    {
        if (current_if->ip == target_ip)
        {
            found_flag = 1;
            target_hw_addr = current_if->addr;
            break;
        }
        current_if = current_if->next;
    }
    
    if (found_flag == 0)
    {
        return;
    }
    else
    {        
        uint8_t *arp_reply_packet = sr_new_ether_arp_pkt(arp_header->ar_sha,
                                                         current_if->addr,
                                                         arp_header->ar_sip,
                                                         current_if->ip,
                                                         0);
        sr_send_packet(sr, arp_reply_packet, ETHER_ARP_PACKET_LEN, current_if->name);
        free(arp_reply_packet);
    }
}

/**
 * Handles an ARP reply. 
 * Looks for corresponding arpreq is ARP request queue and
 *      1. takes the request off the arpreq queue
 *      2. sends the packets waiting for that arpreq
 *      3. inserts an entry in the arpcache.
 */
void sr_handle_arpreply(struct sr_instance *sr, sr_arp_hdr_t *arp_header)
{
    struct sr_arpreq *arpreq = sr_arpcache_insert(&sr->cache, arp_header->ar_sha, 
                                           arp_header->ar_sip);
    struct sr_packet *packet_to_send = arpreq->packets;
    while (packet_to_send) // iterate through linked list of packets
    {
        // Set destination MAC address
        uint8_t *ether_frame = packet_to_send->buf;
        sr_ethernet_hdr_t *ether_hdr = (sr_ethernet_hdr_t *) (ether_frame);
        memcpy(ether_hdr->ether_dhost, arp_header->ar_sha, ETHER_ADDR_LEN);
        
        // TODO must calc new checksum?
        
        sr_send_packet(sr, ether_frame, packet_to_send->len, packet_to_send->iface);
        
        packet_to_send = packet_to_send->next;
    }
    
    sr_arpreq_destroy(&sr->cache, arpreq);    
}



/* Adds an ARP request to the ARP request queue. If the request is already on
   the queue, adds the packet to the linked list of packets for this sr_arpreq
   that corresponds to this ARP request. You should free the passed *packet.
   
   A pointer to the ARP request is returned; it should not be freed. The caller
   can remove the ARP request from the queue by calling sr_arpreq_destroy. */
struct sr_arpreq *sr_arpcache_queuereq(struct sr_arpcache *cache,
                                       uint32_t ip,
                                       uint8_t *packet,           /* borrowed */
                                       unsigned int packet_len,
                                       char *iface)
{
    pthread_mutex_lock(&(cache->lock));
    
    struct sr_arpreq *req;
    for (req = cache->requests; req != NULL; req = req->next) {
        if (req->ip == ip) {
            break;
        }
    }
    
    /* If the IP wasn't found, add it */
    if (!req) {
        req = (struct sr_arpreq *) calloc(1, sizeof(struct sr_arpreq));
        req->ip = ip;
        req->next = cache->requests;
        cache->requests = req;
    }
    
    /* Add the packet to the list of packets for this request */
    if (packet && packet_len && iface) {
        struct sr_packet *new_pkt = (struct sr_packet *)malloc(sizeof(struct sr_packet));
        
        new_pkt->buf = (uint8_t *)malloc(packet_len);
        memcpy(new_pkt->buf, packet, packet_len);
        new_pkt->len = packet_len;
		new_pkt->iface = (char *)malloc(sr_IFACE_NAMELEN);
        strncpy(new_pkt->iface, iface, sr_IFACE_NAMELEN);
        new_pkt->next = req->packets;
        req->packets = new_pkt;
    }
    
    pthread_mutex_unlock(&(cache->lock));
    
    return req;
}

/* This method performs two functions:
   1) Looks up this IP in the request queue. If it is found, returns a pointer
      to the sr_arpreq with this IP. Otherwise, returns NULL.
   2) Inserts this IP to MAC mapping in the cache, and marks it valid. */
struct sr_arpreq *sr_arpcache_insert(struct sr_arpcache *cache, 
                                     unsigned char *mac, 
                                     uint32_t ip)
{
    pthread_mutex_lock(&(cache->lock));
    
    struct sr_arpreq *req, *prev = NULL, *next = NULL; 
    for (req = cache->requests; req != NULL; req = req->next) {
        if (req->ip == ip) {            
            if (prev) {
                next = req->next;
                prev->next = next;
            } 
            else {
                next = req->next;
                cache->requests = next;
            }
            
            break;
        }
        prev = req;
    }
    
    int i;
    for (i = 0; i < SR_ARPCACHE_SZ; i++) {
        if (!(cache->entries[i].valid))
            break;
    }
    
    if (i != SR_ARPCACHE_SZ) {
        memcpy(cache->entries[i].mac, mac, 6);
        cache->entries[i].ip = ip;
        cache->entries[i].added = time(NULL);
        cache->entries[i].valid = 1;
    }
    
    pthread_mutex_unlock(&(cache->lock));
    
    return req;
}

/* Frees all memory associated with this arp request entry. If this arp request 
   entry is on the arp request queue, it is removed from the queue. */
void sr_arpreq_destroy(struct sr_arpcache *cache, struct sr_arpreq *entry) {
    pthread_mutex_lock(&(cache->lock));
    
    if (entry) {
        struct sr_arpreq *req, *prev = NULL, *next = NULL; 
        for (req = cache->requests; req != NULL; req = req->next) {
            if (req == entry) {                
                if (prev) {
                    next = req->next;
                    prev->next = next;
                } 
                else {
                    next = req->next;
                    cache->requests = next;
                }
                
                break;
            }
            prev = req;
        }
        
        struct sr_packet *pkt, *nxt;
        
        for (pkt = entry->packets; pkt; pkt = nxt) {
            nxt = pkt->next;
            if (pkt->buf)
                free(pkt->buf);
            if (pkt->iface)
                free(pkt->iface);
            free(pkt);
        }
        
        free(entry);
    }
    
    pthread_mutex_unlock(&(cache->lock));
}

/* Prints out the ARP table. */
void sr_arpcache_dump(struct sr_arpcache *cache) {
    fprintf(stderr, "\nMAC            IP         ADDED                      VALID\n");
    fprintf(stderr, "-----------------------------------------------------------\n");
    
    int i;
    for (i = 0; i < SR_ARPCACHE_SZ; i++) {
        struct sr_arpentry *cur = &(cache->entries[i]);
        unsigned char *mac = cur->mac;
        fprintf(stderr, "%.1x%.1x%.1x%.1x%.1x%.1x   %.8x   %.24s   %d\n", mac[0], mac[1], mac[2], mac[3], mac[4], mac[5], ntohl(cur->ip), ctime(&(cur->added)), cur->valid);
    }
    
    fprintf(stderr, "\n");
}

/* Initialize table + table lock. Returns 0 on success. */
int sr_arpcache_init(struct sr_arpcache *cache) {  
    /* Seed RNG to kick out a random entry if all entries full. */
    srand(time(NULL));
    
    /* Invalidate all entries */
    memset(cache->entries, 0, sizeof(cache->entries));
    cache->requests = NULL;
    
    /* Acquire mutex lock */
    pthread_mutexattr_init(&(cache->attr));
    pthread_mutexattr_settype(&(cache->attr), PTHREAD_MUTEX_RECURSIVE);
    int success = pthread_mutex_init(&(cache->lock), &(cache->attr));
    
    return success;
}

/* Destroys table + table lock. Returns 0 on success. */
int sr_arpcache_destroy(struct sr_arpcache *cache) {
    return pthread_mutex_destroy(&(cache->lock)) && pthread_mutexattr_destroy(&(cache->attr));
}

/* Thread which sweeps through the cache and invalidates entries that were added
   more than SR_ARPCACHE_TO seconds ago. */
void *sr_arpcache_timeout(void *sr_ptr) {
    struct sr_instance *sr = sr_ptr;
    struct sr_arpcache *cache = &(sr->cache);
    
    while (1) {
        sleep(1.0);
        
        pthread_mutex_lock(&(cache->lock));
    
        time_t curtime = time(NULL);
        
        int i;    
        for (i = 0; i < SR_ARPCACHE_SZ; i++) {
            if ((cache->entries[i].valid) && (difftime(curtime,cache->entries[i].added) > SR_ARPCACHE_TO)) {
                cache->entries[i].valid = 0;
            }
        }
        
        sr_arpcache_sweepreqs(sr);

        pthread_mutex_unlock(&(cache->lock));
    }
    
    return NULL;
}

